#include "armor_detector.hpp"
#include "buff_detector.hpp"
#include "camera.hpp"
#include "compensator.hpp"
#include "spdlog/sinks/basic_file_sink.h"
#include "spdlog/sinks/stdout_color_sinks.h"
#include "spdlog/spdlog.h"

#define AUTO_AIM 1
#define BUFF 0
#define MODULE BUFF

int main(int argc, char const *argv[]) {
  (void)argc;
  (void)argv;

  auto console_sink = std::make_shared<spdlog::sinks::stdout_color_sink_mt>();
  auto file_sink = std::make_shared<spdlog::sinks::basic_file_sink_mt>(
      "logs/ui_params.log", true);

  spdlog::sinks_init_list sink_list = {console_sink, file_sink};

  spdlog::set_default_logger(
      std::make_shared<spdlog::logger>("default", sink_list));

#if (SPDLOG_ACTIVE_LEVEL == SPDLOG_LEVEL_DEBUG)
  spdlog::flush_on(spdlog::level::debug);
  spdlog::set_level(spdlog::level::debug);
#elif (SPDLOG_ACTIVE_LEVEL == SPDLOG_LEVEL_INFO)
  spdlog::flush_on(spdlog::level::info);
  spdlog::set_level(spdlog::level::info);
#endif

  SPDLOG_WARN("***** Running UIParams. *****");
  game::Team team = game::Team::kBLUE;

  Camera cam(0, 640, 480);
  cv::Mat frame, bin_frame;
  const int frame_area = 640 * 480;
#if MODULE == AUTO_AIM
  const std::string window = "window",
                    params_path = "../../../../runtime/RMUL2021_Armor.json";
#else
  const std::string window = "window",
                    params_path = "../../../../runtime/RMUT2021_Buff.json";
  BuffDetector detector(params_path, team);
#endif

  cv::namedWindow(window);
  int params[17] = {0};
  cv::FileStorage fs(params_path,
                     cv::FileStorage::READ | cv::FileStorage::FORMAT_JSON);
  if (fs.isOpened()) {
    params[0] = fs["binary_th"];
    params[1] = fs["se_erosion"];
    params[2] = fs["ap_erosion"];
    params[3] = int(fs["contour_size_low_th"]);
    if (MODULE == AUTO_AIM) {
      params[4] = frame_area * double(fs["contour_area_low_th"]);
      params[5] = frame_area * double(fs["contour_area_high_th"]);
      params[6] = frame_area * double(fs["bar_area_low_th"]);
      params[7] = frame_area * double(fs["bar_area_high_th"]);
      params[8] = 10 * double(fs["angle_high_th"]);
      params[9] = 10 * double(fs["aspect_ratio_low_th"]);
      params[10] = 10 * double(fs["aspect_ratio_high_th"]);

      params[11] = 10 * double(fs["angle_diff_th"]);
      params[12] = 10 * double(fs["length_diff_th"]);
      params[13] = 10 * double(fs["height_diff_th"]);
      params[14] = 10 * double(fs["area_diff_th"]);
      params[15] = 10 * double(fs["center_dist_low_th"]);
      params[16] = 10 * double(fs["center_dist_high_th"]);
    } else if (MODULE == BUFF) {
      params[4] = 10 * double(fs["rect_ratio_low_th"]);
      params[5] = 10 * double(fs["rect_ratio_high_th"]);
      params[6] = fs["contour_center_area_low_th"];
      params[7] = fs["contour_center_area_high_th"];
      params[8] = 10 * double(fs["rect_center_ratio_low_th"]);
      params[9] = 10 * double(fs["rect_center_ratio_high_th"]);
      params[10] = 10 * double(fs["hammer_area_low_th"]);
      params[11] = 10 * double(fs["hammer_area_high_th"]);
      params[12] = 10. * double(fs["armor_area_low_th"]);
      params[13] = 10. * double(fs["armor_area_high_th"]);
    }
    SPDLOG_DEBUG("Loaded ui_params");
  } else {
    SPDLOG_ERROR("Load error! Please initialize JSON file first.");
    return EXIT_FAILURE;
  }
  while (true) {
    frame = cam.GetFrame();
    if (frame.empty()) continue;
    cv::threshold(frame, bin_frame, params[0], 255., cv::THRESH_BINARY);

#if MODULE == AUTO_AIM

    ArmorDetector detector(params_path, team);
    detector.Detect(frame);
    detector.VisualizeResult(frame, 10);

    cv::createTrackbar("binary", window, &params[0], 255);
    cv::createTrackbar("se_erosion", window, &params[1], 5);

    cv::createTrackbar("contour_area_low", window, &params[4],
                       int(0.0001 * frame_area));
    cv::createTrackbar("contour_area_high", window, &params[5],
                       int(0.05 * frame_area));

    cv::createTrackbar("bar_area_low", window, &params[6],
                       int(0.0001 * frame_area));
    cv::createTrackbar("bar_area_high", window, &params[7],
                       int(0.01 * frame_area));
    cv::createTrackbar("angle_high_th", window, &params[8], 900);
    cv::createTrackbar("aspect_ratio_low", window, &params[9], 30);
    cv::createTrackbar("aspect_ratio_high", window, &params[10], 70);
    cv::createTrackbar("angle_diff", window, &params[11], 3000);
    cv::createTrackbar("length_diff", window, &params[12], 50);
    cv::createTrackbar("height_diff", window, &params[13], 50);
    cv::createTrackbar("area_diff", window, &params[14], 100);
    cv::createTrackbar("center_dist_low", window, &params[15], 30);
    cv::createTrackbar("center_dist_high", window, &params[16], 120);

    cv::FileStorage fs(params_path,
                       cv::FileStorage::WRITE | cv::FileStorage::FORMAT_JSON);
    fs << "binary_th" << params[0];
    fs << "se_erosion" << params[1];
    fs << "ap_erosion" << params[2];

    fs << "contour_size_low_th" << params[3];
    fs << "contour_area_low_th" << params[4] * 1. / frame_area;
    fs << "contour_area_high_th" << params[5] * 1. / frame_area;
    fs << "bar_area_low_th" << params[6] * 1. / frame_area;
    fs << "bar_area_high_th" << params[7] * 1. / frame_area;

    fs << "angle_high_th" << params[8] / 10.;
    fs << "aspect_ratio_low_th" << params[9] / 10.;
    fs << "aspect_ratio_high_th" << params[10] / 10.;

    fs << "angle_diff_th" << params[11] / 10.;
    fs << "length_diff_th" << params[12] / 10.;
    fs << "height_diff_th" << params[13] / 10.;
    fs << "area_diff_th" << params[14] / 10.;
    fs << "center_dist_low_th" << params[15] / 10.;
    fs << "center_dist_high_th" << params[16] / 10.;
#else
    detector.Detect(frame);
    detector.VisualizeResult(frame, 3);

    cv::createTrackbar("binary", window, &params[0], 255);
    cv::createTrackbar("se_erosion", window, &params[1], 5);

    cv::createTrackbar("contour_center_area_low", window, &params[6], 1000);
    cv::createTrackbar("contour_center_area_high", window, &params[7], 9000);
    cv::createTrackbar("10rect_center_ratio_low", window, &params[8], 20);
    cv::createTrackbar("10rect_center_ratio_high", window, &params[9], 30);
    cv::createTrackbar("10hammer_area_low", window, &params[10], 100);
    cv::createTrackbar("10hammer_area_high", window, &params[11], 1000);
    cv::createTrackbar("10armor_area_low", window, &params[12], 25);
    cv::createTrackbar("10armor_area_high", window, &params[13], 150);

    cv::FileStorage fs(params_path,
                       cv::FileStorage::WRITE | cv::FileStorage::FORMAT_JSON);
    fs << "binary_th" << params[0];
    fs << "se_erosion" << params[1];
    fs << "ap_erosion" << params[2];

    fs << "contour_size_low_th" << params[3];
    fs << "rect_ratio_low_th" << params[4] / 10.;
    fs << "rect_ratio_high_th" << params[5] / 10.;
    fs << "contour_center_area_low_th" << params[6];
    fs << "contour_center_area_high_th" << params[7];
    fs << "rect_center_ratio_low_th" << params[8] / 10.;
    fs << "rect_center_ratio_high_th" << params[9] / 10.;
    fs << "hammer_area_low_th" << params[10] / 10.;
    fs << "hammer_area_high_th" << params[11] / 10.;
    fs << "armor_area_low_th" << params[12] / 10.;
    fs << "armor_area_high_th" << params[13] / 10.;
#endif

    cv::imshow("bin", bin_frame);
    cv::imshow("show", frame);
    cv::resize(frame, frame, cv::Size(320, 240));
    cv::imshow(window, frame);

    if (cv::waitKey(1) == 'q') {
      return EXIT_SUCCESS;
    }
  }

  return EXIT_SUCCESS;
}
